import os
import cv2


# 定义函数用于读取labels文件并返回目标检测框信息
def read_labels(labels_folder):
    label_info = {}
    for filename in os.listdir(labels_folder):
        if filename.endswith('.txt'):
            image_id = filename.split('.')[0]
            with open(os.path.join(labels_folder, filename), 'r') as file:
                lines = file.readlines()
                objects = []
                for line in lines:
                    object_info = line.strip().split()
                    class_id = int(object_info[0])
                    center_x, center_y = float(object_info[1]), float(object_info[2])
                    width, height = float(object_info[3]), float(object_info[4])
                    x1, y1 = int((center_x - width / 2) * image_width), int((center_y - height / 2) * image_height)
                    x2, y2 = int((center_x + width / 2) * image_width), int((center_y + height / 2) * image_height)
                    objects.append({'class_id': class_id, 'bbox': (x1, y1, x2, y2)})
                label_info[image_id] = objects
    return label_info


# 指定images文件夹和labels文件夹的路径
images_folder = "D:/PycharmProjects/yolo_v8/images/train"
labels_folder = "D:/PycharmProjects/yolo_v8/labels/train"
output_folder = "D:/PycharmProjects/yolo_v8/recognition/images/train"

# 遍历images文件夹中的图像文件
for image_filename in os.listdir(images_folder):
    if image_filename.endswith('.jpg'):
        image_id = image_filename.split('.')[0]
        # 读取图像
        image = cv2.imread(os.path.join(images_folder, image_filename))
        # 检查图像是否成功读取
        if image is None:
            print(f"Error reading image: {image_filename}. Skipping...")
            continue  # 如果图像读取失败，则跳过当前图像，继续下一张图像的处理
        image_height, image_width = image.shape[:2]

        # 读取labels文件
        label_info = read_labels(labels_folder)

        # 检查该图像是否有对应的label信息
        if image_id in label_info:
            objects = label_info[image_id]
            # 遍历图像中的目标检测框信息
            for idx, obj in enumerate(objects):
                x1, y1, x2, y2 = obj['bbox']
                # 在图像上截取目标区域
                object_image = image[y1:y2, x1:x2]
                # 构造输出文件名
                output_filename = f"{image_id}_{idx}.jpg"
                # 保存截取的目标图像
                cv2.imwrite(os.path.join(output_folder, output_filename), object_image)
